package com.example.demo.ai.llm.service;

import com.jfinal.kit.Kv;
import java.util.ArrayList;
import java.util.List;
import org.noear.solon.ai.chat.ChatConfig;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatResponse;
import org.noear.solon.ai.chat.message.ChatMessage;
import org.noear.solon.annotation.Component;
import org.noear.solon.flow.ChainContext;
import org.noear.solon.flow.Node;
import org.noear.solon.flow.TaskComponent;

/**
 * @author airhead
 */
@Component("suggestionTask")
public class SuggestionTask implements TaskComponent {
  @Override
  public void run(ChainContext context, Node node) throws Throwable {
    Kv meta = Kv.create().set(node.meta());
    ChatConfig chatConfig = new ChatConfig();
    chatConfig.setModel(meta.getStr("model"));
    chatConfig.setProvider(meta.getStr("provider"));
    chatConfig.setApiUrl(meta.getStr("apiUrl"));
    ChatModel chatModel = ChatModel.of(chatConfig).build();

    List<ChatMessage> chatMessageList = new ArrayList<>();
    ChatMessage system = ChatMessage.ofSystem(meta.getStr("system"));
    chatMessageList.add(system);

    Kv model = Kv.create().set(context.model());
    String inputKey = meta.getStr("input");
    ChatMessage userMessage = ChatMessage.ofUser(model.getStr(inputKey));
    chatMessageList.add(userMessage);
    ChatResponse response = chatModel.prompt(chatMessageList).call();
    String content = response.getMessage().getContent();

    String outputKey = meta.getStr("output");
    context.put(outputKey, content);

    Kv result = context.get("result");
    result.set("data", content);
  }
}
